# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')

FLAGS = flags.FLAGS


def simple_stats_and_print(merged_df, full_symbol):
  percentiles = [.6, .7, .8, .9]
  print(full_symbol)
  merged_df['time_diff'] = merged_df['timestamp_x'] - merged_df['timestamp_y']
  print(merged_df['time_diff'].describe(percentiles=percentiles))


def drop_dup_and_merge(trade_df1, trade_df2):
  compare_cols = ['price', 'qty', 'side']
  # trade_no_dup1 = trade_df1.drop_duplicates(subset=compare_cols, keep='first')
  # trade_no_dup2 = trade_df2.drop_duplicates(subset=compare_cols, keep='first')
  trade_no_dup1 = trade_df1.groupby(compare_cols, as_index=False).first().reset_index(drop=True)
  trade_no_dup2 = trade_df2.groupby(compare_cols, as_index=False).first().reset_index(drop=True)
  merged_df = pandas.merge(trade_no_dup1, trade_no_dup2, on=compare_cols, how='inner')
  return merged_df


def read_csv_into_df(csv_path):
  df = pandas.read_csv(csv_path, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_dir1 = FLAGS.csv_dir1
  assert csv_dir1, '--csv_dir1 must be specified.'
  csv_dir2 = FLAGS.csv_dir2
  assert csv_dir2, '--csv_dir2 must be specified.'

  for file_name in os.listdir(csv_dir1):
    full_symbol = '.'.join(file_name.split('.')[:-1])
    csv_path1 = os.path.join(csv_dir1, file_name)
    csv_path2 = os.path.join(csv_dir2, file_name)

    trade_df1 = read_csv_into_df(csv_path1)
    trade_df2 = read_csv_into_df(csv_path2)

    merged_df = drop_dup_and_merge(trade_df1, trade_df2)
    simple_stats_and_print(merged_df, full_symbol)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_dir1', None, 'Intput dir1.')

  flags.DEFINE_string('csv_dir2', None, 'Intput dir2.')

  app.run(main)
